package me.mingshan.leetcode;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * https://leetcode.cn/problems/count-nodes-equal-to-average-of-subtree/
 *
 * 给你一棵二叉树的根节点 root ，找出并返回满足要求的节点数，要求节点的值等于其 子树 中值的 平均值 。
 *
 * 注意：
 *
 * n 个元素的平均值可以由 n 个元素 求和 然后再除以 n ，并 向下舍入 到最近的整数。
 * root 的 子树 由 root 和它的所有后代组成。
 *
 * @author hanjuntao
 * @date 2025/10/8 0008
 */
public class L_2256_统计值等于子树平均值的节点数 {

    public static void main(String[] args) {
        TreeNode root = new TreeNode(4);
        TreeNode node1 = new TreeNode(8);
        TreeNode node2 = new TreeNode(0);
        TreeNode node3 = new TreeNode(1);
        TreeNode node4 = new TreeNode(5);
        TreeNode node5 = new TreeNode(6);

        root.left = node1;
        root.left.left = node2;
        root.left.right = node3;

        root.right = node4;
        root.right.right = node5;

        System.out.println(averageOfSubtree(root));
    }

    public static int averageOfSubtree(TreeNode root) {
        AtomicInteger ai = new AtomicInteger();
        dfs(root, ai);

        return ai.get();
    }

    /**
     * 递归
     *
     * 思路：
     * 1. 递归遍历二叉树，返回子树的节点数和子树的和
     * 2. 递归返回结果，并判断当前节点的值是否等于子树平均值，如果是，则计数器加1
     * 3. 递归边界：节点为空，则返回空结果；如果是叶子节点，默认计数器加1
     *
     * @param root
     * @return 包含两部分：子树的节点数，子树的和
     */
    public static Result dfs(TreeNode root, AtomicInteger ai) {
        if (root == null) {
            return new Result();
        }

        if (root.left == null && root.right == null) {
            ai.addAndGet(1);
            return new Result(1, root.val);
        }

        Result leftValue = dfs(root.left, ai);
        Result rightValue = dfs(root.right, ai);

        Result result = new Result();
        result.count = leftValue.count + rightValue.count + 1;
        result.sum = leftValue.sum + rightValue.sum + root.val;

        int avg = result.sum / result.count;
        if (avg == root.val) {
            ai.addAndGet(1);
        }

        return result;
    }

    static class Result {
        int count;
        int sum;

        public Result() {
        }

        public Result(int count, int sum) {
            this.count = count;
            this.sum = sum;
        }
    }
}
